import os
import cv2
import easyocr
import re
import traceback
from PIL import Image

# Pillow兼容性补丁
if not hasattr(Image, 'ANTIALIAS'):
    Image.ANTIALIAS = Image.LANCZOS


class IDCardOCR:
    def __init__(self, lang='ch_sim', use_gpu=False):
        self.reader = easyocr.Reader([lang], gpu=use_gpu, download_enabled=True)

    def recognize(self, image_path, card_side):
        # print(f"正在识别: {image_path}, 卡片面: {card_side}")
        if not os.path.exists(image_path):
            return {"error": "文件不存在"}
        try:
            processed_img = self._preprocess_image(image_path)
            result = self.reader.readtext(processed_img, detail=1)

            # 打印原始OCR结果
            # print("原始OCR结果:", result)

            return self._extract_id_card_info(result, card_side)
        except Exception as e:
            print(f"异常堆栈: {traceback.format_exc()}")
            return {"error": f"识别出错: {str(e)}"}

    def _preprocess_image(self, image_path):
        img = cv2.imread(image_path)
        if img is None:
            raise ValueError(f"无法加载图像: {image_path}")

        # 转为灰度图并增强对比度
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        enhanced = clahe.apply(gray)

        # 保存预处理后的图像用于调试
        cv2.imwrite("preprocessed.jpg", enhanced)

        return enhanced

    def _clean_text(self, text):
        # 使用正则表达式修复字段名中的任意数量空格
        field_mapping = {
            r'姓\s+名': '姓名',
            r'性\s+别': '性别',
            r'民\s+族': '民族',
            r'出\s+生': '出生',
            r'住\s+址': '住址',
            r'公民\s+身份\s+号码': '公民身份号码'
        }

        for pattern, replacement in field_mapping.items():
            text = re.sub(pattern, replacement, text)

        return text.strip()  # 仅去除首尾空格

    def _extract_id_card_info(self, ocr_result, card_side):
        # 降低置信度阈值，捕获更多可能的文本
        texts = [self._clean_text(text) for (bbox, text, prob) in ocr_result if prob > 0.25]
        combined_text = ' '.join(texts)

        # print("合并后的文本:", combined_text)

        info = {}
        if card_side == 'front':
            # 姓名（支持"姓名:张三"、"姓名 张三"、"姓名   张三"等）
            name_match = re.search(r'姓名\s*[:：]?\s*(\S+)', combined_text)
            info['name'] = name_match.group(1) if name_match else None

            # 性别
            gender_match = re.search(r'性别\s*[:：]?\s*(\S)', combined_text)
            info['gender'] = gender_match.group(1) if gender_match else None

            # 民族
            ethnic_match = re.search(r'民族\s*[:：]?\s*(\S+)', combined_text)
            info['ethnicity'] = ethnic_match.group(1) if ethnic_match else None

            # 出生日期
            birth_match = re.search(r'出生\s*[:：]?\s*(\d{4})[年/-]?(\d{1,2})[月/-]?(\d{1,2})日?', combined_text)
            if birth_match:
                year, month, day = birth_match.groups()
                info['birth_date'] = f"{year}年{month}月{day}日"

            # 地址
            address_match = re.search(r'住址\s*[:：]?\s*([\s\S]+?)(公民身份号码|$)', combined_text)
            info['address'] = address_match.group(1).strip() if address_match else None

            # 身份证号码
            id_match = re.search(r'公民身份号码\s*[:：]?\s*(\d{17}[\dXx])', combined_text)
            info['id_number'] = id_match.group(1) if id_match else None

        else:  # 背面信息
            # 签发机关
            issuer_match = re.search(r'签发机关\s*[:：]?\s*([\u4e00-\u9fa5]+)', combined_text)
            info['issuing_authority'] = issuer_match.group(1) if issuer_match else None

            # 有效期限
            validity_match = re.search(r'有效期限\s*[:：]?\s*(\d{4}年\d{1,2}月\d{1,2}日至\d{4}年\d{1,2}月\d{1,2}日)',
                                       combined_text)
            info['valid_period'] = validity_match.group(1) if validity_match else None
        # 返回格式化后的字符串，每行一个字段
        return info